from pyhanlp import *
import jieba.posseg


question_dir = "./newdict/question_classification.txt"
movie_dir = "./newdict/movieDict.txt"
type_dir = "./newdict/typeDict.txt"
director_dir = "./newdict/directorDict.txt"
screenwriter_dir = "./newdict/screenwriterDict.txt"
actor_dir = "./newdict/actorDict.txt"

with open(question_dir, encoding='utf-8') as f:
    question_dir = f.read().splitlines()
    for i, line in enumerate(question_dir):
        question_dir[i] = line.split(':')[-1]

with open(movie_dir, encoding='utf-8-sig') as f:
    movie_dir = f.read().splitlines()
    for i, line in enumerate(movie_dir):
        movie_dir[i] = line

with open(type_dir, encoding='utf-8-sig') as f:
    type_dir = f.read().splitlines()
    for i, line in enumerate(type_dir):
        type_dir[i] = line

with open(director_dir, encoding='utf-8-sig') as f:
    director_dir = f.read().splitlines()
    for i, line in enumerate(director_dir):
        director_dir[i] = line

with open(screenwriter_dir, encoding='utf-8-sig') as f:
    screenwriter_dir = f.read().splitlines()
    for i, line in enumerate(screenwriter_dir):
        screenwriter_dir[i] = line

with open(actor_dir, encoding='utf-8-sig') as f:
    actor_dir = f.read().splitlines()
    for i, line in enumerate(actor_dir):
        actor_dir[i] = line


def dosegment_all(query):
    outstr = ''
    outch = ''
    for x in HanLP.segment(query):
        if x.word in movie_dir:
            outstr += "{},".format(x.word)
            outch += "{},".format('nm')
        elif x.word in type_dir:
            outstr += "{},".format(x.word)
            outch += "{},".format('ng')
        elif x.word in director_dir:
            outstr += "{},".format(x.word)
            outch += "{},".format('nr')
        elif x.word in screenwriter_dir:
            outstr += "{},".format(x.word)
            outch += "{},".format('nr')
        elif x.word in actor_dir:
            outstr += "{},".format(x.word)
            outch += "{},".format('nr')
        else:
            outstr += "{},".format(x.word)
            outch += "{},".format(x.nature)
    outstr = outstr.split(',')
    outch = outch.split(',')
    return outstr, outch


def query_abstract(query):
    str, charact = dosegment_all(query)
    abstract_query = {}
    nrCount = 0
    for ch in charact:
        if ch == 'nm':
            index = charact.index(ch)
            abstract_query[ch] = str[index]
        elif ch == 'nr' and nrCount == 0:
            index = charact.index(ch)
            abstract_query[ch] = str[index]
            charact[index] = '0'
            nrCount += 1
        elif ch == 'nr' and nrCount == 1:
            new_word = 'nnr'
            index = charact.index(ch)
            abstract_query[new_word] = str[index]
        elif ch == 'ng':
            index = charact.index(ch)
            abstract_query[ch] = str[index]
    return abstract_query


def match_question(prediction, query):
    query = query_abstract(query)
    question = question_dir[prediction]
    if 'nm' in question and 'nm' in query.keys():
        question = question.replace('nm', query['nm'])
    if 'nnr' in question and 'nnr' in query.keys():
        question = question.replace('nnr', query['nnr'])
    if 'nr' in question and 'nr' in query.keys():
        question = question.replace('nr', query['nr'])
    if 'ng' in question and 'ng' in query.keys():
        question = question.replace('ng', query['ng'])
    return question


if __name__ == '__main__':
    question = dosegment_all("陈凯歌在哪个地区")
    print(question)
